"""
Created by Kostas Triaridis (@kostino)
in August 2023 @ ITI-CERTH
"""
from collections import defaultdict
from typing import Dict, Tuple, List

import torch
from torch.utils.data import Dataset
import random
import os
import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
from contextlib import contextmanager
import albumentations as albu
from albumentations.core.transforms_interface import DualTransform
from PIL import Image
from torchvision import transforms
from matplotlib import pyplot as plt
from torch.nn import functional as F

import clusterings
import torchvision.transforms.functional as TF
from torchvision.transforms.functional import normalize, resize, to_tensor
from clusterings.misc import get_model, to_one_hot


class EdgeGenerator(torch.nn.Module):
    """generate the 'edge bar' for a 0-1 mask Groundtruth of a image
    Algorithm is based on 'Morphological Dilation and Difference Reduction'

    Which implemented with fixed-weight Convolution layer with weight matrix looks like a cross,
    for example, if kernel size is 3, the weight matrix is:
        [[0, 1, 0],
        [1, 1, 1],
        [0, 1, 0]]

    """

    def __init__(self, kernel_size=3) -> None:
        super().__init__()
        self.kernel_size = kernel_size

    def _dilate(self, image, kernel_size=3):
        """Doings dilation on the image

        Args:
            image (_type_): 0-1 tensor in shape (B, C, H, W)
        """
        assert kernel_size % 2 == 1, "Kernel size must be odd"
        assert image.shape[2] > kernel_size and image.shape[3] > kernel_size, "Image must be larger than kernel size"

        kernel = torch.zeros((1, 1, kernel_size, kernel_size))
        kernel[0, 0, kernel_size // 2: kernel_size // 2 + 1, :] = 1
        kernel[0, 0, :, kernel_size // 2: kernel_size // 2 + 1] = 1
        kernel = kernel.float()
        # print(kernel)
        res = F.conv2d(image, kernel.view([1, 1, kernel_size, kernel_size]), stride=1, padding=kernel_size // 2)
        return (res > 0) * 1.0

    def _find_edge(self, image, kernel_size=3, return_all=False):
        """Find 0-1 edges of the image

        Args:
            image (_type_): 0-1 ndarray in shape (B, C, H, W)
        """
        image = torch.tensor(image).float()
        shape = image.shape

        if len(shape) == 2:
            image = image.reshape([1, 1, shape[0], shape[1]])
        if len(shape) == 3:
            image = image.reshape([1, shape[0], shape[1], shape[2]])
        assert image.shape[1] == 1, "Image must be single channel"

        img = self._dilate(image, kernel_size=kernel_size)

        erosion = self._dilate(1 - image, kernel_size=kernel_size)

        diff = -torch.abs(erosion - img) + 1
        diff = (diff > 0) * 1.0
        # res = dilate(diff)
        diff = diff.numpy()
        if return_all:
            return diff, img, erosion
        else:
            return diff

    def forward(self, x, return_all=False):
        """
        Args:
            image (_type_): 0-1 ndarray in shape (B, C, H, W)
        """
        return self._find_edge(x, self.kernel_size, return_all=return_all)


class RandomCopyMove(DualTransform):
    def __init__(self,
                 max_h=0.8,
                 max_w=0.8,
                 min_h=0.05,
                 min_w=0.05,
                 mask_value=255,
                 always_apply=False,
                 p=0.5,
                 ):
        """Apply cope-move manipulation to the image, and change the respective region on the mask to <mask_value>

        Args:
            max_h (float, optional): (0~1), max window height rate to the full height of image . Defaults to 0.5.
            max_w (float, optional): (0~1), max window width rate to the full width of image . Defaults to 0.5.
            min_h (float, optional): (0~1), min window height rate to the full height of image . Defaults to 0.05.
            min_w (float, optional): (0~1), min window width rate to the full width of image . Defaults to 0.05.
            mask_value (int, optional): the value apply the tampered region on the mask. Defaults to 255.
            always_apply (bool, optional): _description_. Defaults to False.
            p (float, optional): _description_. Defaults to 0.5.
        """
        super(RandomCopyMove, self).__init__(always_apply, p)
        self.max_h = max_h
        self.max_w = max_w
        self.min_h = min_h
        self.min_w = min_w
        self.mask_value = mask_value

    def _get_random_window(
            self,
            img_height,
            img_width,
            window_height=None,
            window_width=None
    ):
        assert self.max_h < 1 and self.max_h > 0
        assert self.max_w < 1 and self.max_w > 0
        assert self.min_w < 1 and self.min_w > 0
        assert self.min_h < 1 and self.min_h > 0

        l_min_h = int(img_height * self.min_h)
        l_min_w = int(img_width * self.min_w)
        l_max_h = int(img_height * self.max_h)
        l_max_w = int(img_width * self.max_w)

        if window_width == None or window_height == None:
            window_h = np.random.randint(l_min_h, l_max_h)
            window_w = np.random.randint(l_min_w, l_max_w)
        else:
            window_h = window_height
            window_w = window_width

        # position of left up corner of the window
        pos_h = np.random.randint(0, img_height - window_h)
        pos_w = np.random.randint(0, img_width - window_w)

        return pos_h, pos_w, window_h, window_w

    def apply(self, img: np.ndarray, **params) -> np.ndarray:
        image = img.copy()
        H, W, _ = image.shape
        # copy region:
        c_pos_h, c_pos_w, c_window_h, c_window_w = self._get_random_window(H, W)

        # past region, window size is defined by copy region:
        self.p_pos_h, self.p_pos_w, self.p_window_h, self.p_window_w = self._get_random_window(H, W, c_window_h,
                                                                                               c_window_w)

        copy_region = image[
                      c_pos_h: c_pos_h + c_window_h,
                      c_pos_w: c_pos_w + c_window_w,
                      :
                      ]
        image[
        self.p_pos_h: self.p_pos_h + self.p_window_h,
        self.p_pos_w: self.p_pos_w + self.p_window_w,
        :
        ] = copy_region
        return image

    def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
        """
        change the mask of manipulated region to 1
        """

        manipulated_region = np.full((self.p_window_h, self.p_window_w), 1)
        img = img.copy()
        img[
        self.p_pos_h: self.p_pos_h + self.p_window_h,
        self.p_pos_w: self.p_pos_w + self.p_window_w,
        ] = self.mask_value
        return img


class RandomInpainting(DualTransform):
    def __init__(self,
                 max_h=0.8,
                 max_w=0.8,
                 min_h=0.05,
                 min_w=0.05,
                 mask_value=255,
                 always_apply=False,
                 p=0.5,
                 ):
        super(RandomInpainting, self).__init__(always_apply, p)
        self.max_h = max_h
        self.max_w = max_w
        self.min_h = min_h
        self.min_w = min_w
        self.mask_value = mask_value

    def _get_random_window(
            self,
            img_height,
            img_width,
    ):
        assert self.max_h < 1 and self.max_h > 0
        assert self.max_w < 1 and self.max_w > 0
        assert self.min_w < 1 and self.min_w > 0
        assert self.min_h < 1 and self.min_h > 0

        l_min_h = int(img_height * self.min_h)
        l_min_w = int(img_width * self.min_w)
        l_max_h = int(img_height * self.max_h)
        l_max_w = int(img_width * self.max_w)

        window_h = np.random.randint(l_min_h, l_max_h)
        window_w = np.random.randint(l_min_w, l_max_w)

        # position of left up corner of the window
        pos_h = np.random.randint(0, img_height - window_h)
        pos_w = np.random.randint(0, img_width - window_w)

        return pos_h, pos_w, window_h, window_w

    def apply(self, img: np.ndarray, **params) -> np.ndarray:
        img = img.copy()
        img = np.uint8(img)
        H, W, C = img.shape
        mask = np.zeros((H, W), dtype=np.uint8)
        # inpainting region
        self.pos_h, self.pos_w, self.window_h, self.window_w = self._get_random_window(H, W)
        mask[
        self.pos_h: self.pos_h + self.window_h,
        self.pos_w: self.pos_w + self.window_w,
        ] = 1
        inpaint_flag = cv2.INPAINT_TELEA if random.random() > 0.5 else cv2.INPAINT_NS
        img = cv2.inpaint(img, mask, 3, inpaint_flag)
        return img

    def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
        """
        change the mask of manipulated region to 1
        """
        img = img.copy()
        img[
        self.pos_h: self.pos_h + self.window_h,
        self.pos_w: self.pos_w + self.window_w,
        ] = self.mask_value
        return img


def get_albu_transforms(type_='train', outputsize=1024):
    """get albumentations transforms

        type_ (str):
            if 'train', then return train transforms with
                random scale, flip, rotate, brightness, contrast, and GaussianBlur augmentation.
            if 'test' then return test transforms
            if 'pad' then return zero-padding transforms
    """

    assert type_ in ['train', 'test', 'pad'], "type_ must be 'train' or 'test' of 'pad' "
    trans = None
    if type_ == 'train':
        trans = albu.Compose([
            # Rescale the input image by a random factor between 0.8 and 1.2
            albu.RandomScale(scale_limit=0.2, p=1),
            RandomCopyMove(p=0.1),
            RandomInpainting(p=0.1),
            # Flips
            albu.Resize(512, 512),
            albu.HorizontalFlip(p=0.5),
            albu.VerticalFlip(p=0.5),
            # Brightness and contrast fluctuation
            albu.RandomBrightnessContrast(
                brightness_limit=(-0.1, 0.1),
                contrast_limit=0.1,
                p=1
            ),
            albu.ImageCompression(
                quality_lower=70,
                quality_upper=100,
                p=0.2
            ),
            # Rotate
            albu.RandomRotate90(p=0.5),
            # Blur
            albu.GaussianBlur(
                blur_limit=(3, 7),
                p=0.2
            ),
        ])

    if type_ == 'test':
        trans = None
        trans = albu.Compose([
            # ---Blow for robustness evalution---
            # albu.Resize(512, 512),
            #   albu.JpegCompression(
            #         quality_lower = 89,
            #         quality_upper = 90,
            #         p = 1
            #   ),
            #  albu.GaussianBlur(
            #         blur_limit = (5, 5),
            #         p = 1
            #     ),

            # albu.GaussNoise(
            #     var_limit=(15, 15),
            #     p = 1
            # )
        ])

    if type_ == 'pad':
        trans = albu.Compose([
            albu.PadIfNeeded(
                min_height=outputsize,
                min_width=outputsize,
                border_mode=0,
                value=0,
                position='top_left',
                mask_value=0),
            albu.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            albu.Crop(0, 0, outputsize, outputsize),
            ToTensorV2(transpose_mask=True)
        ])
    return trans
def pil_loader(path: str) -> Image.Image:
    """PIL image loader

    Args:
        path (str): image path

    Returns:
        Image.Image: PIL image (after np.array(x) becomes [0,255] int8)
    """
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')
@contextmanager
def cwd(path):
    oldpwd = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(oldpwd)


def get_random_crop_coords_on_grid(height: int, width: int, crop_height: int, crop_width: int, h_start: float,
                                   w_start: float):
    y1 = int((height - crop_height + 1) * h_start)
    # MOVE CROP ON JPEG GRID
    y1 = y1 // 8 * 8

    y2 = y1 + crop_height

    x1 = int((width - crop_width + 1) * w_start)
    # MOVE CROP ON JPEG GRID
    x1 = x1 // 8 * 8

    x2 = x1 + crop_width
    return x1, y1, x2, y2


def random_crop_on_grid(img: np.ndarray, crop_height: int, crop_width: int, h_start: float, w_start: float):
    height, width = img.shape[:2]
    if height < crop_height or width < crop_width:
        raise ValueError(
            "Requested crop size ({crop_height}, {crop_width}) is "
            "larger than the image size ({height}, {width})".format(
                crop_height=crop_height, crop_width=crop_width, height=height, width=width
            )
        )
    x1, y1, x2, y2 = get_random_crop_coords_on_grid(height, width, crop_height, crop_width, h_start, w_start)
    img = img[y1:y2, x1:x2]
    return img


class RandomCropONJPEGGRID(A.RandomCrop):
    def __init__(self, height, width, always_apply=False, p=1.0):
        super().__init__(height, width, always_apply, p)
        self.height = height
        self.width = width

    def apply(self, img, h_start=0, w_start=0, **params):
        return random_crop_on_grid(img, self.height, self.width, h_start, w_start)

    def get_params(self):
        return {"h_start": random.random(), "w_start": random.random()}


def get_next_power_of_2(x):
    return 1 << (x - 1).bit_length()


class ManipulationDataset(Dataset):
    def __init__(self, path, image_size, train=False):
        self.train = train
        self.path = path
        self.image_size = image_size
        self.image_paths = []
        self.mask_paths = []
        self.base_path = './data'
        self.labels = []
        self.name = self.path.split('/')[-1].replace('.txt', '').replace('IDT-', '')
        self.h = 0
        self.w = 0
        self.pad_test = 0
        #----------------#
        # IML边缘提取器设置
        #----------------#
        # self.train_transform = get_albu_transforms('train')
        self.edge_generator = EdgeGenerator(7)
        self.edge_generator.to('cuda:2')
        # self.padding_transform = get_albu_transforms('pad',outputsize=512)
        # ----------------#
        # fore
        # ----------------#
        # load a model
        # device: torch.device = torch.device("cuda:0")
        # self.network = get_model(
        #     arch={"mocov2": "resnet50", "swav": "resnet50", "dino": "vit_small"}["dino"],
        #     training_method="dino",
        #     patch_size=16  # for dino
        # )
        # self.network.eval()

        self._init_transforms()
        with open(self.path, 'r') as f:
            lines = f.readlines()
            for line in lines:
                parts = line.rstrip().split(' ')
                try:
                    image_path, mask_path, label_str = parts
                except ValueError:
                    print("Incorrect amount of columns in file, expected 4")
                    print(parts)

                self.image_paths.append(image_path)
                self.mask_paths.append(mask_path)
                self.labels.append(int(label_str))

    def __len__(self):
        return len(self.labels)

    def __repr__(self):
        return self.name

    def _init_transforms(self):

        self.image_transforms_train = A.Compose([

            ##IML-ViT###
            # Rescale the input image by a random factor between 0.8 and 1.2
            # albu.RandomScale(scale_limit=0.2, p=1),
            RandomCopyMove(p=0.1),
            RandomInpainting(p=0.1),
            # Flips
            # albu.Resize(512, 512),
            albu.HorizontalFlip(p=0.5),
            albu.VerticalFlip(p=0.5),
            # Brightness and contrast fluctuation
            albu.RandomBrightnessContrast(
                brightness_limit=(-0.1, 0.1),
                contrast_limit=0.1,
                p=1
            ),
            # albu.ImageCompression(
            #     quality_lower=70,
            #     quality_upper=100,
            #     p=0.2
            # ),
            # Rotate
            albu.RandomRotate90(p=0.5),
            # Blur
            albu.GaussianBlur(
                blur_limit=(3, 7),
                p=0.2
            ),

            A.RandomScale(scale_limit=(-0.5, 0.5), p=0.5),
            A.PadIfNeeded(min_height=self.image_size, min_width=self.image_size, border_mode=cv2.BORDER_CONSTANT,
                          value=127, mask_value=-1, p=1),
            A.RandomCrop(height=self.image_size, width=self.image_size, p=1),
            A.ImageCompression(quality_lower=30, quality_upper=100, p=0.5),

        ])

        self.image_transforms_resize = A.Compose([

            # Rescale the input image by a random factor between 0.8 and 1.2
            # albu.RandomScale(scale_limit=0.2, p=1),
            # RandomCopyMove(p=0.1),
            # RandomInpainting(p=0.1),
            # # Flips
            # # albu.Resize(512, 512),
            # albu.HorizontalFlip(p=0.5),
            # albu.VerticalFlip(p=0.5),
            # # Brightness and contrast fluctuation
            # albu.RandomBrightnessContrast(
            #     brightness_limit=(-0.1, 0.1),
            #     contrast_limit=0.1,
            #     p=1
            # ),
            # # albu.ImageCompression(
            # #     quality_lower=70,
            # #     quality_upper=100,
            # #     p=0.2
            # # ),
            # # Rotate
            # albu.RandomRotate90(p=0.5),
            # # Blur
            # albu.GaussianBlur(
            #     blur_limit=(3, 7),
            #     p=0.2
            # ),
            ########--------------------#########
            A.RandomScale(scale_limit=(-0.5, 0.5), p=0.5),
            A.PadIfNeeded(min_height=self.image_size, min_width=self.image_size, border_mode=cv2.BORDER_CONSTANT,
                          value=127, mask_value=-1, p=1),
            A.RandomCrop(height=self.image_size, width=self.image_size, p=1),
            # A.Resize(height=self.image_size,width=self.image_size,p=1),
            A.ImageCompression(quality_lower=30, quality_upper=100, p=0.5),
        ])

        self.image_transforms_final = A.Compose([
            # 鲁棒性实验
            # A.GaussianBlur(blur_limit=(23, 23), p=1.0),  # 高斯模糊 (3, 5, 7, 9, 11, 15,19,23)
            # A.ImageCompression(quality_lower=50, quality_upper=50,p = 1),  # jpeg压缩 (100, 90, 80, 70, 60, 50)
            # A.RandomGamma(gamma_limit=(70,70),p=1),  # Gamma  [130, 115, 100, 85, 70]
            # A.GaussNoise(var_limit=(529,529), p=1),  # [9, 49 , 121, 225, 361, 529]
            ToTensorV2()
        ])

        self.trans = albu.Compose([
            # albu.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2(transpose_mask=True)
        ])

    # def mask_to_bbox(self, mask: np.ndarray) -> Dict[int, Tuple[int, int, int, int]]:
    #     """Given a binary mask, return a list of bounding box coordinates (ymin, ymax, xmin, xmax)."""
    #     mask_index_to_bbox = dict()
    #     if len(mask.shape) == 2:
    #         mask = mask[None]
    #
    #     for mask_index, m in enumerate(mask):
    #         y_coords, x_coords = np.where(m)
    #         try:
    #             ymin, ymax, xmin, xmax = np.min(y_coords), np.max(y_coords), np.min(x_coords), np.max(x_coords)
    #         except ValueError:  # a mask which does not predict anything.
    #             continue
    #         mask_index_to_bbox[mask_index] = (ymin.item(), ymax.item(), xmin.item(), xmax.item())
    #     return mask_index_to_bbox
    #
    # def filter_masks(
    #         self,
    #         dt_masks: torch.Tensor,
    #         mask_index_to_bbox: dict,
    #         remove_long_masks: bool = True,
    #         remove_small_large_masks: bool = False
    # ) -> Tuple[torch.Tensor, Dict[int, int]]:
    #     list_filtered_masks: list = list()
    #     new_index_to_prev_index: dict = dict()
    #     h, w = dt_masks.shape[-2:]
    #     new_index = 0
    #
    #     for mask_index, bbox in mask_index_to_bbox.items():
    #         ymin, ymax, xmin, xmax = bbox
    #         if remove_long_masks:
    #             if ymin == 0 and ymax + 1 == h:
    #                 continue
    #             elif xmin == 0 and xmax + 1 == w:
    #                 continue
    #
    #         if remove_small_large_masks:
    #             if dt_masks[mask_index].sum() < 0.05 * h * w:
    #                 continue
    #             elif (xmax - xmin) * (ymax - ymin) > 0.95 * h * w:
    #                 continue
    #         list_filtered_masks.append(dt_masks[mask_index])
    #         new_index_to_prev_index[new_index] = mask_index
    #         new_index += 1
    #     try:
    #         return torch.stack(list_filtered_masks, dim=0), new_index_to_prev_index
    #     except RuntimeError:  # rare case where all predictions are filtered.
    #         return dt_masks, {i: i for i in range(len(dt_masks))}
    #
    # def vote_mask(
    #         self,
    #         batch_pred_masks,
    #         remove_long_masks: bool = True,
    #         remove_small_large_masks: bool = False
    # ):
    #     mask_index_to_bbox = self.mask_to_bbox(batch_pred_masks.squeeze(dim=0).cpu().numpy())
    #     masks, new_index_to_prev_index = self.filter_masks(
    #         dt_masks=batch_pred_masks.squeeze(dim=0),
    #         mask_index_to_bbox=mask_index_to_bbox,
    #         remove_long_masks=remove_long_masks,
    #         remove_small_large_masks=remove_small_large_masks
    #     )  # n x h x w
    #     masks_t = masks.permute(1, 2, 0)
    #     try:
    #         intersection = torch.logical_and(masks[..., None], masks_t[None]).sum(dim=(1, 2))  # n x n
    #         union = torch.logical_or(masks[..., None], masks_t[None]).sum(dim=(1, 2))  # n x n
    #     except RuntimeError:
    #         masks = masks.cpu()
    #         masks_t = masks_t.cpu()
    #         intersection = torch.logical_and(masks[..., None], masks_t[None]).sum(dim=(1, 2))  # n x n
    #         union = torch.logical_or(masks[..., None], masks_t[None]).sum(dim=(1, 2))  # n x n
    #
    #     iou_table = intersection / (union + 1e-7)  # n x n
    #     ious = iou_table.sum(dim=1)  # n
    #     sorted_index = torch.argsort(ious, descending=True)
    #     best_mask_index = sorted_index[0].cpu().item()
    #     best_mask = masks[best_mask_index]
    #     return best_mask, best_mask_index, new_index_to_prev_index
    #
    # def pad_input_image(self, image: torch.Tensor, total_stride: int):
    #     assert len(image.shape) == 4
    #     h_image, w_image = image.shape[-2:]
    #     pad_w = (total_stride - w_image % total_stride) % total_stride
    #     pad_h = (total_stride - h_image % total_stride) % total_stride
    #     image: torch.Tensor = TF.pad(image, [0, 0, pad_w, pad_h])
    #
    #     h_padded_image, w_padded_image = h_image + pad_h, w_image + pad_w
    #     h_feat, w_feat = h_padded_image // total_stride, w_padded_image // total_stride
    #     return image, h_feat, w_feat
    #
    # def extract_candidate_masks(self, batch_imgs, image_path, feature_types) -> Dict[str, np.ndarray]:
    #     filename_to_candidate_masks: Dict[str, list] = defaultdict(list)
    #
    #     cluster_sizes: Tuple[int, ...] = (2, 3, 4)
    #     use_gpu: bool = True  # whether to use a gpu for clustering
    #     cluster_type: str = "spectral"
    #     if cluster_type == "k-means":
    #         clusterer: callable = clusterings.KMeansClustering(use_gpu=use_gpu)
    #     else:
    #         clusterer: callable = clusterings.SpectralClustering(use_gpu=use_gpu)
    #     # feature_types: List[str] = ["mocov2", "swav", "dino"],
    #     feature_types: List[str] = feature_types
    #
    #
    #     for feature_type in feature_types:
    #         # print(f"========== Generating candidate masks with {feature_type} ==========")
    #         # load a model
    #         # network = get_model(
    #         #     arch={"mocov2": "resnet50", "swav": "resnet50", "dino": "vit_small"}[feature_type],
    #         #     training_method=feature_type,
    #         #     patch_size=16  # for dino
    #         # )
    #         # network.eval()
    #
    #         # img: Image.Image = Image.open(p_images).convert("RGB")
    #         # # print(img.size)
    #         # # image = denormalize(image)
    #         # # 定义转换
    #         # # transform = transforms.ToTensor()
    #         # #
    #         # # # 转换为 Tensor
    #         # # tensor = transform(img)
    #         # # image = tensor.squeeze(0).permute(1, 2, 0)
    #         # # image = image.cpu().numpy()
    #         # # plt.figure(figsize=(5, 3))
    #         # # # plt.imshow(salient_mask.cpu().numpy())
    #         # # plt.imshow(image)
    #         # # plt.axis("off")
    #         # # plt.tight_layout(pad=0)
    #         # # plt.show()
    #         # # plt.close()
    #         #
    #         # mean: Tuple[float, float, float] = (0.485, 0.456, 0.406)
    #         # std: Tuple[float, float, float] = (0.229, 0.224, 0.225)
    #         # img = normalize(to_tensor(img).unsqueeze(0), mean=mean, std=std)
    #
    #         # batch_imgs: torch.Tensor = img.to(device)  # b (=1) x 3 x H x W
    #         h_image, w_image = batch_imgs.shape[-2:]
    #
    #         filename: str = image_path.split("/")[-1]
    #         # extract features froom a given model
    #         if feature_type in ["mocov2", "swav"]:
    #             # dilated resnet50
    #             total_stride: int = 8
    #             batch_imgs, h_feat, w_feat = self.pad_input_image(batch_imgs, total_stride=total_stride)
    #             features: torch.Tensor = self.network(batch_imgs)[-1]
    #         else:
    #             # ViT-S/16
    #             total_stride: int = 16
    #             batch_imgs, h_feat, w_feat = self.pad_input_image(batch_imgs, total_stride=total_stride)
    #             try:
    #                 batch_tokens = self.network(batch_imgs, layer="layer12")
    #             except AttributeError:
    #                 batch_tokens = self.network.encoder(batch_imgs, layer="layer12")
    #
    #             batch_patch_tokens = batch_tokens[:, 1:, :]  # b (=1) x (h_feat * w_feat) x n_dims
    #
    #             # batch_patch_tokens
    #             batch_patch_tokens = batch_patch_tokens.view(batch_imgs.shape[0], h_feat, w_feat, -1)
    #             features = batch_patch_tokens.permute(0, 3, 1, 2)
    #
    #         # upsample by 2 before clustering
    #         features = F.interpolate(features, scale_factor=2, mode="bilinear", align_corners=True)
    #
    #         # iterate over a list of given cluster sizes
    #         for k in cluster_sizes:
    #             # clustering
    #             # batch_clusters: b (=1) x h x w
    #             batch_clusters: np.ndarray = clusterer(features, k)
    #
    #             # batch_one_hot_masks: b (=1) x k x h x w -> k x h x w
    #             batch_one_hot_masks: np.ndarray = to_one_hot(batch_clusters[0])
    #
    #             batch_one_hot_masks: torch.Tensor = F.interpolate(
    #                 batch_one_hot_masks[None],
    #                 scale_factor=(total_stride // 2, total_stride // 2),
    #                 mode="nearest"
    #             )[0]
    #
    #             batch_one_hot_masks: torch.Tensor = batch_one_hot_masks[..., :h_image, :w_image]
    #             one_hot_masks: np.ndarray = batch_one_hot_masks.numpy().astype(np.uint8)
    #             filename_to_candidate_masks[filename].append(one_hot_masks)
    #
    #     # concatenate a list of masks for each image into a numpy array
    #     for filename, candidate_mask in filename_to_candidate_masks.items():
    #         filename_to_candidate_masks[filename]: np.ndaray = np.concatenate(candidate_mask, axis=0)
    #
    #     return filename_to_candidate_masks

    def __getitem__(self, index):
        # --------------------
        # Read image and label
        # --------------------
        # with cwd(self.base_path):
        try:
            image = cv2.cvtColor(cv2.imread(self.image_paths[index]), cv2.COLOR_BGR2RGB)
        except Exception as e:
            print(e)
            print(self.image_paths[index])
            print(image.shape)
        h, w, c = image.shape
        label = self.labels[index]

        # ----------
        # Read mask
        # ----------
        if self.mask_paths[index] == 'None':
            mask = np.zeros((h, w), np.uint8)  # a totally black mask for real image
        else:
            # with cwd(self.base_path):
            if self.mask_paths[index].endswith('.npz'):
                data = np.load(self.mask_paths[index])
                mask = data['arr_0'].squeeze()  # numpy
            else:
                mask = cv2.imread(self.mask_paths[index], cv2.IMREAD_GRAYSCALE)
        h, w, _ = image.shape
        h_m, w_m = mask.shape
        if h != h_m or w != w_m:
            resize = transforms.Resize([h, w])
            mask = Image.fromarray(mask)  # numpy -> PIL对象
            mask = resize(mask)
            mask = np.array(mask)  # numpy

        # if self.train:
        #     # #-----------------
        #     # # read orig_img
        #     # #-----------------
        #     # if self.name is "CASIA_v2_manip":
        #     #     ori = self.image_paths[index].split('_')[-3]
        #     #     ori_path = "/raid/datasets/ImageForgery/CASIAv2/Au/" + 'Au_' + ori[0:3] + '_' + ori[3:] + '.jpg'
        #     #     orig_img = cv2.cvtColor(cv2.imread(ori_path), cv2.COLOR_BGR2RGB)
        #     # elif self.name is "FantasticReality_manip":
        #     #     ori = self.image_paths[index].split('_')
        #     #     ori_path = "/raid/datasets/ImageForgery/FantasticReality_v1/dataset/ColorRealImages/" + ori[-2] +'_' + ori[-1]
        #     #     orig_img = cv2.cvtColor(cv2.imread(ori_path), cv2.COLOR_BGR2RGB)
        #     # elif self.name is "cm_COCO":
        #     #     ori = self.image_paths[index].split('/')[-1].split('.jpg')[0].split('_')[1]
        #     #     ori_path = "/raid/datasets/PoseEstimation/coco/train2017/" + ori + ".jpg"
        #     #     orig_img = cv2.cvtColor(cv2.imread(ori_path), cv2.COLOR_BGR2RGB)
        #     # elif self.name is "sp_COCO":
        #     #     ori = self.image_paths[index].split('/')[-1].split(".jpg")[0].split('_')[1] + ".jpg"
        #     #     ori_path = "/raid/datasets/PoseEstimation/coco/train2017/" + ori
        #     #     orig_img = cv2.cvtColor(cv2.imread(ori_path), cv2.COLOR_BGR2RGB)
        #     # else:
        #     #     orig_img = image    # numpy (h,w,c) uint8
        #     # print(ori_path)
        #     # ------------------
        #     # Read fore
        #     # ------------------
        #     filename: str = self.image_paths[index].split("/")[-1]
        #     image_path = os.path.join("/home/wc/disk1/datasets/foreground", self.name, filename + "_fore.png")
        #     fore = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        #     fore = fore / 255
        #     fore = fore.astype(np.int8)   # numpy (h,w)  int8
        # else:
        #     # ori_img = np.random.randint(0, 256, size=(512, 512, 3), dtype=np.uint8)
        #     candidate_masks = np.random.randint(0, 1, size=(512, 512), dtype=np.int8)
        # if self.train:
        #     resize_img_list = ["FantasticReality_auth", "FantasticReality_manip", "IMD2020", "compRAISE"]
        #     if self.name in resize_img_list:
        #         res = self.image_transforms_resize(image=image, mask=mask)
        #         image = res['image']
        #         mask = res['mask']
        #         # res = self.image_transforms_resize(image=orig_img)
        #         # ori_img = res['image']
        #         candidate_masks = fore
        #     else:
        #         mask_list = [mask,fore]
        #         res = self.image_transforms_train(image=image, masks=mask_list)
        #         image = res['image']
        #         mask = res['masks'][0]
        #         candidate_masks = res['masks'][1]
        #         # res = self.image_transforms_train(mask=fore)
        #         # ori_img = res['image']
        #         # candidate_masks = res['mask']
        if self.train:
            res = self.image_transforms_train(image=image, mask=mask)
            image = res['image']
            mask = res['mask']
        elif h > 1024 or w > 1024:
            res = A.LongestMaxSize(max_size=1024)(image=image, mask=mask)
            image = res['image']
            mask = res['mask']
        # if self.train is False:
        #     h = get_next_power_of_2(image.shape[0])
        #     w = get_next_power_of_2(image.shape[1])
        #
        #     image_transforms_val = A.Compose([
        #         A.PadIfNeeded(min_height=h, min_width=w, border_mode=cv2.BORDER_CONSTANT,
        #                       value=127, mask_value=-1, p=1),
        #     ])
        #     res = image_transforms_val(image=image,mask=mask)
        #     image = res['image']
        #     mask = res['mask']

        # # 保存图像
        # outer_folder_name = "foreground"
        # folder_name = self.name
        # filename: str = self.image_paths[index].split("/")[-1]
        # image_path = os.path.join(outer_folder_name, folder_name, filename + "_fore.png")
        # # 检查文件夹是否存在，如果不存在则创建
        # os.makedirs(os.path.join(outer_folder_name, folder_name), exist_ok=True)
        # if not os.path.exists(image_path):
        #     mean: Tuple[float, float, float] = (0.485, 0.456, 0.406)
        #     std: Tuple[float, float, float] = (0.229, 0.224, 0.225)
        #     img = normalize(to_tensor(image).unsqueeze(0), mean=mean, std=std)
        #     # print(img.shape)  # torch.Size([1, 3, 512, 512])
        #     masks = self.extract_candidate_masks(img, self.image_paths[index], ['dino'])
        #     for filename, candidate_masks in masks.items():
        #         candidate_masks: torch.Tensor = torch.tensor(candidate_masks)[None]  # 1 x k x h x w
        #
        #         # salient_mask: torch.Tensor, h x w
        #         salient_mask, _, _ = self.vote_mask(
        #             candidate_masks,
        #             remove_long_masks=True,
        #             remove_small_large_masks=False
        #         )
        #         # filename_to_encoded_salient_mask[filename] = encode(np.asfortranarray(salient_mask.cpu().numpy()))
        #         candidate_masks = salient_mask.cpu().numpy()
        #         # print(f"未保存candidate_masks.shape:{candidate_masks.shape}...{type(candidate_masks)}")
        #         # print(candidate_masks)
        #         # print(candidate_masks.shape)  # (512, 512)
        #     data = candidate_masks * 255
        #     data = data.astype(np.uint8)  # 转换为无符号整数类型
        #     # 将 NumPy 数组转换为 PIL 图像
        #     fore = Image.fromarray(data)
        #     fore.save(image_path)
        #     print(f"保存图片: {image_path}")
        # else:
        #     print(f"图片 {filename} 已经存在于 {folder_name} 中")
        #     candidate_masks = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        #     candidate_masks = candidate_masks / 255
        #     candidate_masks = candidate_masks.astype(np.int8)
        # mask_list.append(candidate_masks)

        # else:
        #     res = self.image_transforms_val(image = image,masks=mask_list)
        #     image = res['image']
        #     mask = res['mask'][0]
        #     fore = res['mask'][1]

            # print(f"已保存candidate_masks.shape:{candidate_masks.shape}...{type(candidate_masks)}")
            # print(candidate_masks)

            # print(candidate_masks.shape)

            # text_file_path = "fore_name.txt"
            # # 以追加模式打开文本文件并写入图片文件名
            # with open(text_file_path, 'a') as f:  # 使用 'a' 模式追加写入
            #     f.write(image_path + '\n')  # 写入文件名，并换行

            # print(f"创建文件夹: {outer_folder_name}/{folder_name}")
            # if not os.path.exists(folder_name):
            #     os.makedirs(folder_name)
            #     print(f"创建文件夹: {folder_name}")
        image = self.image_transforms_final(image=image)['image']
        # ori_image = self.image_transforms_final(image=ori_img)['image']
        # print(image.shape)  # torch.Size([3, 512, 512])
        image = image / 256.0  # [c,h,w]
        # ori_image = ori_image / 256.0  # [c,h,w]
        #--------------
        #  Read edge
        #--------------
        gt_img = (mask > 127.5) * 1.0
        edge = self.edge_generator(gt_img)
        res = self.trans(image=edge[0][0])
        edge = res['image'].squeeze(0)
        # print(edge.shape)
        # print(edge.shape)  #[1,1,h,w]
        # plt.subplot(1, 2, 1)
        # plt.imshow(edge[0, 0, :, :], cmap='gray')
        # plt.subplot(1, 2, 2)
        # plt.imshow(gt_img[:, :], cmap='gray')
        # plt.show()
        mask = mask / 255.0
        mask = ToTensorV2()(image=mask)['image']
        mask = (mask > 0.1).long()  # [h,w]
        # mask = 1.0 - mask
        # candidate_masks = candidate_masks * 1.0
        # candidate_masks = torch.tensor(candidate_masks)
        # candidate_masks = candidate_masks.unsqueeze(0)
        # print(f'Images shape: {image.shape}, type: {type(image)}')
        # print(f'masks shape: {mask.shape}, type: {type(mask)}')
        # print(f'fore shape: {candidate_masks.shape}, type: {type(candidate_masks)}')
        # print(f"candidate_masks.shape:{candidate_masks.shape}")  #torch.Size([1, 512, 512])
        # return image, [], mask, label,edge, candidate_masks, _

        # samples = {}
        # samples['img'] = image
        # samples['mask'] = mask
        # return samples
        return image, [], mask, label,edge


class MixDataset(Dataset):
    def __init__(self, paths, image_size, train=False, class_weight=None):
        self.train = train
        self.paths = paths
        self.image_size = image_size
        self.dataset_list = []
        for path in paths:
            self.dataset_list.append(ManipulationDataset(path, image_size, train=train))

        if class_weight is None:
            self.class_weights = torch.FloatTensor([1.0, 1.0])
        else:
            self.class_weights = torch.FloatTensor(class_weight)

        self.lens = [len(d) for d in self.dataset_list]

        self.smallest = min(self.lens)

    # Should be shuffled every epoch to make sure different samples from big datasets are selected
    def shuffle(self):
        for dat in self.dataset_list:
            temp = list(zip(dat.image_paths, dat.mask_paths, dat.labels))
            random.shuffle(temp)
            image_paths, mask_paths, labels = zip(*temp)
            dat.image_paths, dat.mask_paths, dat.labels = list(image_paths), list(mask_paths), list(labels)

    def __len__(self):
        return len(self.dataset_list) * self.smallest

    def get_info(self):
        s = "Using datasets:\n"
        for ds in self.dataset_list:
            s += (str(ds) + ' - ' + str(len(ds)) + ' Images')
            s += '\n'
        s += "Data Configuration: "
        s += f"crop_size={self.image_size}, class_weight={self.class_weights}\n"
        return s

    def __getitem__(self, index):
        # balanced sampling
        if index < self.smallest * len(self.dataset_list):
            return self.dataset_list[index // self.smallest][index % self.smallest]
        else:
            raise ValueError("Something wrong.")


from torch.utils.data import DataLoader
#
if __name__ == '__main__':
    config = {
        'DATASET': [
            '/home/wc/disk1/MMFusion/data/CAT-Net_splits/train/CASIA_v2_auth.txt',
            '/home/wc/disk1/MMFusion/data/CAT-Net_splits/train/CASIA_v2_manip.txt'
        ]
    }
    train = MixDataset(config['DATASET'],
                       512,
                       train=True,
                       class_weight=[0.5, 2.5])
    train_loader = DataLoader(train,
                              batch_size=4,
                              shuffle=True,
                              num_workers=16,
                              pin_memory=True)
    iters_per_epoch = len(train_loader)
    print(iters_per_epoch)
    for step, (images, _, masks, _,_,_,_) in enumerate(train_loader):
        print("=")
